//===--- OwnedToGuaranteedTransform.cpp -----------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "fso-owned-to-guaranteed-transform"

#include "polarphp/pil/optimizer/internal/funcsignaturetransforms/FunctionSignatureOpts.h"
#include "polarphp/pil/lang/DebugUtils.h"
#include "llvm/Support/CommandLine.h"

using namespace polar;

static llvm::cl::opt<bool> FSODisableOwnedToGuaranteed(
   "sil-fso-disable-owned-to-guaranteed",
   llvm::cl::desc("Do not perform owned to guaranteed during FSO. Intended "
                  "only for testing purposes."));

//===----------------------------------------------------------------------===//
//                                 Utilities
//===----------------------------------------------------------------------===//

/// Return the single return value of the function.
static PILValue findReturnValue(PILFunction *F) {
   auto RBB = F->findReturnBB();
   if (RBB == F->end())
      return PILValue();
   auto Term = dyn_cast<ReturnInst>(RBB->getTerminator());
   return Term->getOperand();
}

/// Return the single apply found in this function.
static PILInstruction *findOnlyApply(PILFunction *F) {
   PILInstruction *OnlyApply = nullptr;
   for (auto &B : *F) {
      for (auto &X : B) {
         if (!isa<ApplyInst>(X) && !isa<TryApplyInst>(X))
            continue;
         assert(!OnlyApply && "There are more than 1 function calls");
         OnlyApply = &X;
      }
   }
   assert(OnlyApply && "There is no function calls");
   return OnlyApply;
}

//===----------------------------------------------------------------------===//
//                               Implementation
//===----------------------------------------------------------------------===//

bool FunctionSignatureTransform::OwnedToGuaranteedAnalyzeParameters() {
   PILFunction *F = TransformDescriptor.OriginalFunction;
   auto Args = F->begin()->getPILFunctionArguments();
   // A map from consumed PILArguments to the release associated with an
   // argument.
   //
   // TODO: The return block and throw block should really be abstracted away.
   PILArgumentConvention ArgumentConventions[] = {
      PILArgumentConvention::Direct_Owned, PILArgumentConvention::Indirect_In};
   ConsumedArgToEpilogueReleaseMatcher ArgToReturnReleaseMap(
      RCIA->get(F), F, ArgumentConventions);
   ConsumedArgToEpilogueReleaseMatcher ArgToThrowReleaseMap(
      RCIA->get(F), F, ArgumentConventions,
      ConsumedArgToEpilogueReleaseMatcher::ExitKind::Throw);

   // Did we decide we should optimize any parameter?
   bool SignatureOptimize = false;

   // Analyze the argument information.
   for (unsigned i : indices(Args)) {
      ArgumentDescriptor &A = TransformDescriptor.ArgumentDescList[i];
      if (!A.canOptimizeLiveArg()) {
         continue;
      }

      // See if we can find a ref count equivalent strong_release or release_value
      // at the end of this function if our argument is an @owned parameter.
      // See if we can find a destroy_addr at the end of this function if our
      // argument is an @in parameter.
      if (A.hasConvention(PILArgumentConvention::Direct_Owned) ||
          A.hasConvention(PILArgumentConvention::Indirect_In)) {
         auto Releases = ArgToReturnReleaseMap.getReleasesForArgument(A.Arg);
         if (!Releases.empty()) {
            // If the function has a throw block we must also find a matching
            // release in the throw block.
            auto ReleasesInThrow =
               ArgToThrowReleaseMap.getReleasesForArgument(A.Arg);
            if (!ArgToThrowReleaseMap.hasBlock() || !ReleasesInThrow.empty()) {
               assert(A.CalleeRelease.empty());
               assert(A.CalleeReleaseInThrowBlock.empty());
               llvm::copy(Releases, std::back_inserter(A.CalleeRelease));
               llvm::copy(ReleasesInThrow,
                          std::back_inserter(A.CalleeReleaseInThrowBlock));
               // We can convert this parameter to a @guaranteed.
               A.OwnedToGuaranteed = true;
               SignatureOptimize = true;
            }
         }
      }

      // Modified self argument.
      if (A.OwnedToGuaranteed && Args[i]->isSelf()) {
         TransformDescriptor.shouldModifySelfArgument = true;
      }
   }
   return SignatureOptimize;
}

bool FunctionSignatureTransform::OwnedToGuaranteedAnalyzeResults() {
   PILFunction *F = TransformDescriptor.OriginalFunction;
   auto ResultDescList = TransformDescriptor.ResultDescList;

   auto fnConv = F->getConventions();
   // For now, only do anything if there's a single direct result.
   if (fnConv.getNumDirectPILResults() != 1)
      return false;
   if (!fnConv.getIndirectPILResults().empty())
      return false;

   bool SignatureOptimize = false;
   if (ResultDescList[0].hasConvention(ResultConvention::Owned)) {
      auto RV = findReturnValue(F);
      if (!RV)
         return false;
      auto &RI = ResultDescList[0];
      // We have an @owned return value, find the epilogue retains now.
      auto Retains = EA->get(F)->computeEpilogueARCInstructions(
         EpilogueARCContext::EpilogueARCKind::Retain, RV);
      // We do not need to worry about the throw block, as the return value is
      // only going to be used in the return block/normal block of the try_apply
      // instruction.
      if (!Retains.empty()) {
         RI.CalleeRetain = Retains;
         SignatureOptimize = true;
         RI.OwnedToGuaranteed = true;
      }
   }
   return SignatureOptimize;
}

void
FunctionSignatureTransform::OwnedToGuaranteedTransformFunctionParameters() {
   // And remove all Callee releases that we found and made redundant via owned
   // to guaranteed conversion.
   for (const ArgumentDescriptor &AD : TransformDescriptor.ArgumentDescList) {
      if (!AD.OwnedToGuaranteed)
         continue;
      for (auto &X : AD.CalleeRelease) {
         X->eraseFromParent();
      }
      for (auto &X : AD.CalleeReleaseInThrowBlock) {
         X->eraseFromParent();
      }

      // Now we need to replace the FunctionArgument so that we have the correct
      // ValueOwnershipKind.
      AD.Arg->setOwnershipKind(ValueOwnershipKind::Guaranteed);
   }
}

void FunctionSignatureTransform::OwnedToGuaranteedTransformFunctionResults() {
   // And remove all callee retains that we found and made redundant via owned
   // to unowned conversion.
   for (const ResultDescriptor &RD : TransformDescriptor.ResultDescList) {
      if (!RD.OwnedToGuaranteed)
         continue;
      for (auto &X : RD.CalleeRetain) {
         if (isa<StrongRetainInst>(X) || isa<RetainValueInst>(X)) {
            X->eraseFromParent();
            continue;
         }
         // Create a release to balance it out.
         auto AI = cast<ApplyInst>(X);
         createDecrementBefore(AI, AI->getParent()->getTerminator());
      }
   }
}

void FunctionSignatureTransform::OwnedToGuaranteedFinalizeThunkFunction(
   PILBuilder &Builder, PILFunction *F) {
   // Finish the epilogue work for the argument as well as result.
   for (auto &ArgDesc : TransformDescriptor.ArgumentDescList) {
      OwnedToGuaranteedAddArgumentRelease(ArgDesc, Builder, F);
   }
   for (auto &ResDesc : TransformDescriptor.ResultDescList) {
      OwnedToGuaranteedAddResultRelease(ResDesc, Builder, F);
   }
}

static void createArgumentRelease(PILBuilder &Builder, ArgumentDescriptor &AD) {
   auto &F = Builder.getFunction();
   PILArgument *Arg = F.getArguments()[AD.Index];
   if (Arg->getType().isAddress()) {
      assert(AD.PInfo->getConvention() == ParameterConvention::Indirect_In &&
             F.getConventions().useLoweredAddresses());
      Builder.createDestroyAddr(RegularLocation::getAutoGeneratedLocation(),
                                F.getArguments()[AD.Index]);
      return;
   }
   Builder.createReleaseValue(RegularLocation::getAutoGeneratedLocation(),
                              F.getArguments()[AD.Index],
                              Builder.getDefaultAtomicity());
}

/// Set up epilogue work for the thunk arguments based in the given argument.
/// Default implementation simply passes it through.
void FunctionSignatureTransform::OwnedToGuaranteedAddArgumentRelease(
   ArgumentDescriptor &AD, PILBuilder &Builder, PILFunction *F) {
   // If we have any arguments that were consumed but are now guaranteed,
   // insert a releasing RC instruction.
   if (!AD.OwnedToGuaranteed) {
      return;
   }

   PILInstruction *Call = findOnlyApply(F);
   if (isa<ApplyInst>(Call)) {
      Builder.setInsertionPoint(&*std::next(PILBasicBlock::iterator(Call)));
      createArgumentRelease(Builder, AD);
   } else {
      PILBasicBlock *NormalBB = dyn_cast<TryApplyInst>(Call)->getNormalBB();
      Builder.setInsertionPoint(&*NormalBB->begin());
      createArgumentRelease(Builder, AD);

      PILBasicBlock *ErrorBB = dyn_cast<TryApplyInst>(Call)->getErrorBB();
      Builder.setInsertionPoint(&*ErrorBB->begin());
      createArgumentRelease(Builder, AD);
   }
}

void FunctionSignatureTransform::OwnedToGuaranteedAddResultRelease(
   ResultDescriptor &RD, PILBuilder &Builder, PILFunction *F) {
   // If we have any result that were consumed but are now guaranteed,
   // insert a releasing RC instruction.
   if (!RD.OwnedToGuaranteed) {
      return;
   }

   PILInstruction *Call = findOnlyApply(F);
   if (auto AI = dyn_cast<ApplyInst>(Call)) {
      Builder.setInsertionPoint(&*std::next(PILBasicBlock::iterator(AI)));
      Builder.createRetainValue(RegularLocation::getAutoGeneratedLocation(), AI,
                                Builder.getDefaultAtomicity());
   } else {
      PILBasicBlock *NormalBB = cast<TryApplyInst>(Call)->getNormalBB();
      Builder.setInsertionPoint(&*NormalBB->begin());
      Builder.createRetainValue(RegularLocation::getAutoGeneratedLocation(),
                                NormalBB->getArgument(0),
                                Builder.getDefaultAtomicity());
   }
}

bool FunctionSignatureTransform::OwnedToGuaranteedAnalyze() {
   if (FSODisableOwnedToGuaranteed)
      return false;

   bool Result = OwnedToGuaranteedAnalyzeResults();
   bool Params = OwnedToGuaranteedAnalyzeParameters();
   return Params || Result;
}

void FunctionSignatureTransform::OwnedToGuaranteedTransform() {
   OwnedToGuaranteedTransformFunctionResults();
   OwnedToGuaranteedTransformFunctionParameters();
}
